#include "state.h"


State::State(int n, int argc, char *argv[])
    : common_dimension(n), Process(argc, argv)
{
#if USE_MPI || USE_SHMEM
    try  {
        if (power_2((long long)size) < 0)
            throw "Error: Use 2^k process";
    } catch (const char* str) {
        if (rank == 0)
            cout << str << endl;
        exit(2);
    }
    dimension = common_dimension - power_2((long long)size);
    common_length = 1L << common_dimension;
    length = 1L  << dimension;
    begin = rank * length;
    end = begin + length - 1;
#elif USE_OPENMP
    length = common_length = 1L << common_dimension;
    dimension = common_dimension;
    begin = 0;
    end = length - 1;
#endif
    define_size();
    vec = new complexd [length];
}


State::State(const State &obj)
{
    dimension = obj.dimension;
    common_dimension = obj.common_dimension;
    length = obj.length;
    common_length = obj.common_length;
    size_block = obj.size_block;
    begin = obj.begin;
    end = obj.end;
    rank = obj.rank;
    size = obj.size;
    vec = new complexd [length];
    for (long long i = 0; i < length; i++)
        vec[i] = obj.vec[i];
}


State::~State()
{
    delete [] vec;
}


void State::doConversion(Conversion &conversion, int k)
{
    complexd *vec_tmp = 0, *vec_new = 0;
    long long l, index, real_index, block;
    int n = common_dimension, k_bit;
#ifdef USE_MPI
    MPI_Status status;
#endif

    complexd t[] = {conversion[0], conversion[1],
                    conversion[2], conversion[3]};

// define rank of communicate processor
    if ((l = set_k_bit_i_0(begin,k,n)) == begin)
         l = set_k_bit_i_1(begin,k,n);
    int proc = what_is_rank(l);

    if (proc != rank)
        vec_tmp = new complexd [size_block];
    else {
        vec_tmp = vec;
        vec_new = new complexd [length]; }

    for (l = 0; l < length / size_block; l++) {
        block = size_block * l;

// send and recieve to buffer vec_tmp
        if (proc != rank)
#ifdef USE_OPENMP
        ;
#elif USE_MPI
            MPI_Sendrecv(&vec[block], size_block, MPI_DOUBLE_COMPLEX, proc, l,
                             vec_tmp, size_block, MPI_DOUBLE_COMPLEX, proc, l,
                                                     MPI_COMM_WORLD, &status);
#elif USE_SHMEM
            ;
#endif
        for (long long i = 0; i < size_block; i++) {
            index = block + i;
            real_index = begin + index;
            k_bit = k_bit_of_i(real_index, k, n);
            if (rank < proc)
                vec[index] = t[k_bit * 2 + 0] * vec[index]
                           + t[k_bit * 2 + 1] * vec_tmp[i];
            else if (rank > proc)
                vec[index] = t[k_bit * 2 + 0] * vec_tmp[i]
                           + t[k_bit * 2 + 1] * vec[index];
            else
                vec_new[index] = t[k_bit * 2 + 0] * vec[set_k_bit_i_0(real_index,k,n) - begin]
                               + t[k_bit * 2 + 1] * vec[set_k_bit_i_1(real_index,k,n) - begin];
        }
    }
    barrierAll();
    if (vec_tmp != vec)
        delete [] vec_tmp;
    if (vec_new != 0) {
        delete [] vec;
        vec = vec_new;
    }
}

void State::nConversion(Conversion &con, double noize)
{
    for (int i = 1; i <= common_dimension; i++) {
        con.doNoize(noize);
        //if (rank == root)
            //cout << con << endl << endl << endl;
        doConversion(con, i);
    }
}

void State::n0Conversion(Conversion &con)
{
    for (int i = 1; i <= common_dimension; i++)
        doConversion(con, i);
}


void State::define_size()
{
    size_block = (length > SIZE_BLOCK ? SIZE_BLOCK : length);
}

int State::what_is_rank(long long index)
{
    return index / length;
}

int State::what_is_index(long long index)
{
    return index % length;
}

int State::power_2(long long base)
{
    int power, i;
    for (i = 1, power = 0; i < base; power++)
        i = i << 1;
    if (i ==  base)
        return power;
    else
        return -1;
}


void State::norma()
{
    long double sum = 0;
    for (long long i = 0; i < length; i++)
        sum += abs(vec[i]) * abs(vec[i]);
#ifdef USE_MPI
    //#MPI_All
    long double* sum_buf;
    if (rank == root)
        sum_buf = new long double [size];
    MPI_Gather(&sum, 1, MPI_LONG_DOUBLE,
            sum_buf, 1, MPI_LONG_DOUBLE, root, MPI_COMM_WORLD);
    if (rank == root) {
        sum = 0;
        for (int i = 0; i < size; i++)
            sum += sum_buf[i];
        sum = sqrt(sum);
    }
    MPI_Bcast(&sum, 1, MPI_LONG_DOUBLE, root, MPI_COMM_WORLD);
#endif
    for (long long i = 0; i < length; i++)
        vec[i] = vec[i] / (double)sum;
}


void State::randomFill()
{
   int timemy = time(NULL);
#ifdef USE_OPENMP
    my_omp_srand(time);
#elif USE_MPI || USE_SHMEM
    srand(timemy * rank);
#endif
    for (long long i = 0; i < length; i++)
        vec[i] = myrandom();
}

void State::intentionalFill(long long com_length, complexd* new_vec)
{
    common_length = com_length;
    try  {
        if ((common_dimension = power_2(common_length)) < 0)
            throw "Error: vector must be 2^k";
    } catch (const char* str) {
        if (rank == 0)
            cout << str << endl;
        exit(3);
    }
    delete [] vec;
    dimension = common_dimension - power_2((long long)size);
    length = 1L << dimension;
    define_size();
    begin = rank * length;
    end = begin + length - 1;
    vec = new complexd [length];
    for (long long i = 0; i < length; i++)
        vec[i] = new_vec[i + begin];
}

State State::operator=(State state)
{
    common_dimension = state.common_dimension;
    dimension = state.dimension;
    common_length = state.common_length;
    length = state.length;
    begin = state.begin;
    end = state.end;
    size_block = state.size_block;
    rank = state.rank;
    size = state.size;
    vec = new complexd [length];
    for (int i = 0; i < length; i++)
        vec[i] = state.vec[i];
    return *this;
}

bool State::operator==(State state)
{
    double eps = 0.001;
    int answer = 1;
    if (common_dimension != state.common_dimension ||
        common_length!= state.common_length        ||
        dimension != state.dimension               ||
        length != state.length
       )
        answer = 0;
    if (answer != 0) {
        for (int i = 0; i < length; i++)
            if (abs(vec[i] != state.vec[i]) > eps) {
                answer = 0;
cout << "THis is i = " << i << endl;
                break;
            }

    }
#ifdef USE_MPI
    int* rbuf;
    if (rank == root)
        rbuf = new int [size];
    MPI_Gather(&answer,1, MPI_INT, rbuf,1, MPI_INT, root, MPI_COMM_WORLD);
    barrierAll();
    if (rank == root)
        for (int i = 0; i < size; i++)
            if (rbuf[i]  != 1) {
                answer = 0;
                break;
            }
    MPI_Bcast(&answer, 1, MPI_INT, root, MPI_COMM_WORLD);
#endif
    return (bool)answer;
}

bool State::operator!=(State state)
{
    if (*this == state)
        return false;
    else
        return true;
}

complexd State::operator*(State state)
{
    complexd sum(0,0);
    for (int i = 0; i < length; i++)
        sum += vec[i] * conj(state.vec[i]);
#ifdef USE_MPI
    complexd *sum_buf;
    if (rank == root)
        sum_buf = new complexd [size];
    MPI_Gather(&sum, 1, MPI_DOUBLE_COMPLEX,
            sum_buf, 1, MPI_DOUBLE_COMPLEX, root, MPI_COMM_WORLD);
    if (rank == root) {
        sum = complexd(0,0);
        for (int i = 0; i < size; i++)
            sum += sum_buf[i];
    }
    MPI_Bcast(&sum, 1, MPI_DOUBLE_COMPLEX, root, MPI_COMM_WORLD);
#endif
    return sum;
}

ostream& operator<<(ostream &stream, State &state)
{
#ifdef USE_MPI || USE_SHMEM
    long long int pointer = 0;
#ifdef USE_MPI
    MPI_Status status;
#endif
    for (int i = 0; i < state.size; i++) {
        if (state.rank == i) {
            if (i != 0)
                #ifdef USE_MPI
                MPI_Recv(&pointer, 1, MPI_LONG_LONG_INT, i - 1, 0, MPI_COMM_WORLD, &status);
                #elif USE_SHMEM
                    ;
                #endif
            stream.seekp(pointer, ios::beg);
            for (long long j = 0; j < state.length; j++)
                stream << state.vec[j] << " ";
            if (i != state.size - 1 && state.size != 1) {
                stream.flush();
                pointer = stream.tellp();
                #ifndef USE_MPI
                MPI_Send(&pointer, 1, MPI_LONG_LONG_INT, i + 1, 0, MPI_COMM_WORLD);
                #elif USE_SHMEM
                    ;
                #endif
            } else {
                stream << endl;
                stream.flush();
            }
        }
    }
    state.barrierAll();
#elif USE_OPENMP
    for (long long i = 0; i < state.length; i++)
        stream << state.vec[i] << " ";
    stream  << endl;
    stream.flush();
#endif
    return stream;
}

istream& operator>>(istream &stream, State &state)
{
#ifdef USE_MPI || USE_SHMEM
    long long int pointer = 0;
#ifdef USE_MPI
    MPI_Status status;
#endif
    for (int i = 0; i < state.size; i++) {
        if (state.rank == i) {
            if (i != 0)
                #ifdef USE_MPI
                MPI_Recv(&pointer, 1, MPI_LONG_LONG_INT, i - 1, 0, MPI_COMM_WORLD, &status);
                #elif USE_SHMEM
                    ;
                #endif
            stream.seekg(pointer, ios::beg);
            for (long long j = 0; j < state.length; j++)
                stream >> state.vec[j];
            if (i != state.size - 1 && state.size != 1) {
                pointer = stream.tellg();
                #ifndef USE_MPI
                MPI_Send(&pointer, 1, MPI_LONG_LONG_INT, i + 1, 0, MPI_COMM_WORLD);
                #elif USE_SHMEM
                    ;
                #endif
            }
        }
    }
    state.barrierAll();
#elif USE_OPENMP
    for (long long i = 0; i < state.length; i++)
        stream << state.vec[i] << " ";
    stream  << endl;
    stream.flush();
#endif
    return stream;
}


void State::output(char *filename)
{
#ifdef USE_MPI
    MPI_File out_fd;
    MPI_File_open(MPI_COMM_WORLD, filename, MPI_MODE_WRONLY | MPI_MODE_CREATE,
                                    MPI_INFO_NULL, &out_fd);
    MPI_File_set_view(out_fd, 0, MPI_DOUBLE_COMPLEX, MPI_DOUBLE_COMPLEX, "native", MPI_INFO_NULL);
    MPI_File_write_at_all(out_fd, rank * length,
                                            vec, length, MPI_DOUBLE_COMPLEX, 0);
    MPI_File_close(&out_fd);
#endif
}

void State::input(char *filename)
{
#ifdef USE_MPI
    MPI_File in_fd;
    MPI_File_open(MPI_COMM_WORLD, filename, MPI_MODE_RDONLY,
                                    MPI_INFO_NULL, &in_fd);
    MPI_File_set_view(in_fd, 0, MPI_DOUBLE_COMPLEX, MPI_DOUBLE_COMPLEX, "native", MPI_INFO_NULL);
    MPI_File_read_at_all(in_fd, rank * length,
                                          vec, length, MPI_DOUBLE_COMPLEX, 0);
    MPI_File_close(&in_fd);
#endif
}


void State::info()
{
    cout << "My rank = " << rank  << "(all rocess " << size << ")"<< endl;

    cout << "Dimension = "        << dimension        << endl;
    cout << "Common dimension = " << common_dimension << endl;
    cout << "Length = "           << length           << endl;
    cout << "Common length = "    << common_length    << endl;
    cout << "Begin = "            << begin            << endl;
    cout << "End = "              << end              << endl;

}


int State::getDimension()
{
    return common_dimension;
}

int State::getMyDimension()
{
    return dimension;
}

long long State::getLong()
{
    return common_length;
}

long long State::getMyLong()
{
    return length;
}

complexd* State::getMyVector()
{
    return vec;
}

